import numpy as np
import copy

class SDDetector(object):

    def __init__(self, params):

        self.Nr = params['Nr']
        self.Nt = params['Nt']
        self.c  = params['CONSTELLATION']
        self.batch_size = params['BATCH_SIZE']
        self.C = np.expand_dims(params['CONSTELLATION_COMPLEX'],0) # [1,M]
        self.g = params['CONSTELLATION'][1] - params['CONSTELLATION'][0]

    def __call__(self,x,H,y,sigma2):
        """
        球译码检测算法，仅适用于 QAM 调制
        Input:
        x:发送符号 array(shape=(batch_size, 2*Nt), dtype=float32)
        y:接收信号 array(shape=(batch_size, 2*Nr), dtype=float32)
        H:信道矩阵 array(shape=(batch_size, 2*Nr, 2*Nt), dtype=float32)
        sigma2:噪声方差 array(shape=(batch_size, 1), dtype=float32)
        Output:
        nodes:指标字典
        """
        # 为了提升运算速度，仅关心球译码算法，使用真实发送符号 x 计算 d
        Hx = np.squeeze(np.matmul(H,np.expand_dims(x,-1)),-1)
        d = np.sqrt(np.sum(np.square(y-Hx),1))

        xhat = np.zeros_like(x)
        for i in range(self.batch_size):
            outr, outx = self.SD(y[i,:],H[i,:,:],d[i]+1e-8)
            minindex = np.argmin(outr)
            xhat[i,:] = outx[minindex]

        xhat_complex = xhat[:,0:self.Nt] + 1j*xhat[:,self.Nt:self.Nt*2]
        x_complex  = x[:,0:self.Nt] + 1j*x[:,self.Nt:self.Nt*2]

        totalnum = self.batch_size * self.Nt
        rightnum = self.equal(xhat_complex, x_complex)
        errornum = totalnum - rightnum

        nodes={
            'x_complex':x_complex,
            'xhat_complex':xhat_complex,
            'totalnum':totalnum,
            'errornum':errornum,
        }

        return nodes

    def SD(self, y, H, d):
        '''
        球译码算法，仅适用于 QAM 调制
        '''

        Q, R = np.linalg.qr(H,'complete')

        Q1, Q2, R = Q[:,:2*self.Nt], Q[:,2*self.Nt:], R[:2*self.Nt,:]
        z = Q1.T @ y
        
        d_ = np.zeros((2*self.Nt+1,))
        x_ = np.zeros((2*self.Nt+1,))
        u_ = np.zeros((2*self.Nt+1,))
        z_ = np.zeros((2*self.Nt+1,))
        outx = []
        outr = []

        k = 2 * self.Nt
        d_[k] = np.sqrt(max(0,d**2 - np.linalg.norm(Q2.T@y)**2))
        z_[k] = z[k-1]
        u_[k] = self.rounds( max( (d_[k]+z_[k])/R[k-1,k-1],(-d_[k]+z_[k])/R[k-1,k-1]) )
        x_[k] = self.roundr( min( (d_[k]+z_[k])/R[k-1,k-1],(-d_[k]+z_[k])/R[k-1,k-1]) ) - self.g

        while k != 2*self.Nt+1 :
            x_[k] += self.g
            if x_[k] <= u_[k] :
                if k == 1 :
                    outx.append(copy.deepcopy(x_[1:]))
                    outr.append( d_[2*self.Nt]**2 - d_[1]**2 + (z_[1]-R[1-1,1-1]*x_[1])**2 )
                else :
                    k -= 1
                    z_[k] = z[k-1] - np.sum(x_[k+1:] * R[k-1,k+1-1:])
                    d_[k] = np.sqrt(max(0,d_[k+1]**2 - (z_[k+1] - R[k+1-1,k+1-1]*x_[k+1])**2))
                    u_[k] = self.rounds(max( (d_[k]+z_[k])/R[k-1,k-1],(-d_[k]+z_[k])/R[k-1,k-1]))
                    x_[k] = self.roundr(min( (d_[k]+z_[k])/R[k-1,k-1],(-d_[k]+z_[k])/R[k-1,k-1])) - self.g
            else :
                k += 1

        if outx :
            return outr, outx
        else :
            return self.SD(y, H, d+0.000001)

    def rounds(self,x):

        for point in self.c[::-1]:
            if point <= x : break

        return point

    def roundr(self,x):

        for point in self.c:
            if point >= x : break

        return point

    def batch_matvec_mul(self, A, b):
        """
        矩阵A与矩阵b相乘，其中A.shape=(batch_size, Nr, Nt)
        b.shape = (batch_size, Nt)
        输出矩阵C，C.shape = (batch_size, Nr)
        """
        C = np.matmul(A, np.expand_dims(b, axis=2))

        return np.squeeze(C, -1)

    def slicer(self, x):
        """
        软符号转换为硬符号
        Input:
        x: 检测器检测后的信号 Tensor(shape=(batchsize, 2*Nt), dtype=float32)
        Output:
        xhat_complex: 符号硬估计，Tensor(shape=(batchsize, Nt)， dtype=tf.int32)
        """
        x_complex = x[:, 0: self.Nt] + 1j*x[:, self.Nt: self.Nt*2]
        x_complex = np.reshape(x_complex,[-1, 1])
        indices = np.argmin(np.abs(x_complex - self.C), axis=1).astype(np.int32)
        xhat_complex = self.C[:,indices].T
        xhat_complex = np.reshape(xhat_complex,[-1,self.Nt])
        
        return xhat_complex

    def equal(self,x,y):
        """
        Computes the numbers of elements for which x and y are equal
        """
        return np.sum(np.equal(x, y).astype(np.float32))